package cn.lg.soar.mvc.filter;

import cn.lg.soar.common.util.data.DataUtil;
import cn.lg.soar.mvc.model.IpFilterProperties;
import cn.lg.soar.mvc.model.IpFilterRule;
import jakarta.servlet.FilterRegistration;
import jakarta.servlet.ServletContext;
import jakarta.servlet.ServletException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.web.servlet.ServletContextInitializer;
import org.springframework.boot.web.servlet.support.SpringBootServletInitializer;
import org.springframework.core.Ordered;

import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * @author luguoxiang
 * @date 2022/9/5
 * 开源项目：https://gitee.com/lgx1992/lg-soar 求star！请给我star！请帮我点个star！
 */
public class IpFilterRegistrar extends SpringBootServletInitializer implements ServletContextInitializer {

    @Autowired
    private IpFilterProperties props;

    @Override
    public void onStartup(ServletContext servletContext) throws ServletException {
        // 白名单
        List<IpFilterRule> whitelist = props.getWhitelist();
        if (DataUtil.isValuable(whitelist)) {
            // 循环增加拦截器
            int index = 0;
            for (IpFilterRule rule : whitelist) {
                String[] paths = resolveUrlPatterns(rule.getPaths());
                Set<String> ips = resolveIps(rule.getIps());
                if (DataUtil.isEmpty(paths) || DataUtil.isEmpty(ips)) {
                    throw new RuntimeException("ip-whitelist configuration error：Must include the “path-patterns” and “ips”");
                }
                IpWhitelistFilter filter = new IpWhitelistFilter(ips);
                // 注册过滤器
                FilterRegistration.Dynamic dynamic = servletContext.addFilter("if-whitelist-filter-" + index++, filter);
                dynamic.addMappingForUrlPatterns(null, false, paths); // 设置过滤所有请求
            }
        }
        // 黑名单
        List<IpFilterRule> blacklist = props.getBlacklist();
        if (DataUtil.isValuable(blacklist)) {
            // 循环增加拦截器
            int order = Ordered.HIGHEST_PRECEDENCE + 1;
            int index = 0;
            for (IpFilterRule rule : blacklist) {
                String[] paths = resolveUrlPatterns(rule.getPaths());
                Set<String> ips = resolveIps(rule.getIps());
                if (DataUtil.isEmpty(paths) || DataUtil.isEmpty(ips)) {
                    throw new RuntimeException("ip-blacklist configuration error：Must include the “path-patterns” and “ips”");
                }
                IpBlacklistFilter filter = new IpBlacklistFilter(ips);
                // 注册过滤器
                FilterRegistration.Dynamic dynamic = servletContext.addFilter("if-blacklist-filter-" + index++, filter);
                dynamic.addMappingForUrlPatterns(null, false, paths); // 设置过滤所有请求
            }
        }
    }

    private Set<String> resolveIps(Set<String> ips) {
        if (ips == null) {
            return null;
        }
        return ips.stream()
                .map(String::trim)
                .filter(s -> !s.isEmpty())
                .collect(Collectors.toSet());
    }

    private String[] resolveUrlPatterns(Set<String> paths) {
        if (paths == null) {
            return null;
        }
        return paths.stream()
                .map(String::trim)
                .filter(s -> !s.isEmpty())
                .toArray(String[]::new);
    }
}